# -*- coding: utf-8 -*-
# @Author : Gan
# @Time : 2024/6/20 9:24

import os
import csv
from tqdm import tqdm
from glob import glob
import numpy as np
from concurrent.futures import ProcessPoolExecutor


def load_reference_images(reference_path):
    references = {}
    for path in glob(os.path.join(reference_path, '*')):
        try:
            base_img_num, y1_pred, x1_pred, y2_pred, x2_pred = path.split('\\')[-1][:-4].split('_')
            idx = path.split('\\')[-1].split('_')[0]
            references.setdefault(idx, []).append((np.array([int(x1_pred), int(y1_pred), int(x2_pred), int(y2_pred)]), path.split('\\')[-1]))
        except ValueError:
            print(f"Warning: Skipping file {path}, filename format is incorrect.")
    return references

def calculate_ious(gt_box, pred_boxes):
    """
    计算一个真实框与多个预测框之间的交并比 (IoU)。
    参数:
    - gt_box: 形状为(1, 4)的numpy数组,表示真实框的左上、右下像素坐标。
    - pred_boxes: 形状为(n, 4)的numpy数组,表示n个预测框的左上、右下像素坐标。
    返回:
    - 一个长度为n的numpy数组,包含每个预测框与真实框的交并比。
    """
    # 提取坐标
    xmin_gt, ymin_gt, xmax_gt, ymax_gt = gt_box.squeeze()
    xmin_pred, ymin_pred, xmax_pred, ymax_pred = pred_boxes.T
    # 计算交集的坐标
    inter_xmin = np.maximum(xmin_gt, xmin_pred)
    inter_ymin = np.maximum(ymin_gt, ymin_pred)
    inter_xmax = np.minimum(xmax_gt, xmax_pred)
    inter_ymax = np.minimum(ymax_gt, ymax_pred)
    # 计算交集面积，避免负数
    inter_area = np.maximum(0, inter_xmax - inter_xmin) * np.maximum(0, inter_ymax - inter_ymin)
    # 计算真实框和预测框各自的面积
    area_gt = (xmax_gt - xmin_gt) * (ymax_gt - ymin_gt)
    area_pred = (xmax_pred - xmin_pred) * (ymax_pred - ymin_pred)
    # 计算并集面积
    union_area = area_gt + area_pred - inter_area
    # 计算IoU
    iou = inter_area / union_area
    # 处理可能出现的除以零情况（当预测框或真实框面积为0时）
    iou = np.where(union_area == 0, 0, iou)  # 如果并集面积为0，则IoU设为0
    return iou

def process_query_image(query_name, reference_images):
    idx, y1_gt, x1_gt, y2_gt, x2_gt = query_name[:-4].split('_')
    gt = np.array([int(x1_gt), int(y1_gt), int(x2_gt), int(y2_gt)]).reshape(1, 4)
    if idx in reference_images:
        # Correctly extract coordinate arrays from tuples
        pred_boxes = np.array([coords for coords, _ in reference_images[idx]])
        ious = calculate_ious(gt, pred_boxes)
        match_idx = np.argmax(ious)
        match_coords, reference_name = reference_images[idx][match_idx]
        return query_name, match_coords, np.max(ious), reference_name
    else:
        return query_name, None, None, None

if __name__ == '__main__':
    query_path = r'E:\datasets\satgeoloc_dataset/test'
    reference_path = r'E:\datasets\satgeoloc_dataset/reference'

    query_list = os.listdir(query_path)
    reference_images = load_reference_images(reference_path)

    results = []
    for img_name in tqdm(query_list):
        result = process_query_image(img_name, reference_images)
        results.append(result)

    # 优化：直接将结果写入CSV，避免在内存中构建大型数据结构
    with open(r"E:\datasets\satgeoloc_dataset\match_data_test.csv", "w", encoding="utf-8", newline="") as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(
            ['query_name', 'x1_gt', 'y1_gt', 'x2_gt', 'y2_gt', 'center_gt_x', 'center_gt_y', 'reference_name', 'x1_pred', 'y1_pred', 'x2_pred',
             'y2_pred', 'center_pred_x', 'center_pred_y', 'max_iou'])

        for query_name, match_box, max_iou, reference_name in results:
            if match_box is not None:
                # 从img_name中提取GT坐标
                idx, y1_gt, x1_gt, y2_gt, x2_gt = query_name[:-4].split('_')
                # 从match_box中提取预测坐标
                x1_pred, y1_pred, x2_pred, y2_pred = match_box
                writer.writerow(
                    [query_name, y1_gt, x1_gt,  y2_gt, x2_gt, (int(y1_gt) + int(y2_gt)) / 2, (int(x1_gt) + int(x2_gt)) / 2,  reference_name, y1_pred, x1_pred, y2_pred, x2_pred,
                     (int(y1_pred) + int(y2_pred)) / 2, (int(x1_pred) + int(x2_pred)) / 2,  f"{max_iou:.4f}"])